from flask import Flask, request
from flask_cors import CORS
import base64
from io import BytesIO
import cv2
import numpy as np
from corners import detect_corners

app = Flask(__name__)
CORS(app)

def cv_engine(img, operation):
    if operation == 'to_grayscale':
        return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    elif operation == 'get_edge_canny':
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        return cv2.Canny(gray, 100, 200)
    elif operation == 'get_corners':
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        return detect_corners(gray)
    else:
        raise ValueError("Unsupported operation")

def read_image(image_data):
    image_data = base64.decodebytes(image_data)
    img_array = np.frombuffer(image_data, dtype=np.uint8)
    img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError("Image decoding failed")
    return img

def encode_image(img):
    _, buffer = cv2.imencode('.jpg', img)
    return base64.b64encode(buffer).decode('utf-8')

@app.route('/process_image', methods=['POST'])
def process_image():
    try:
        if not request.json or 'msg' not in request.json:
            return 'Server Error!', 500

        image_data = request.json['image_data'][23:].encode()
        operation = request.json['operation']
        img = read_image(image_data)
        img_out = cv_engine(img, operation)
        image_data = encode_image(img_out)
        return image_data, 200
    except Exception as e:
        print("Error:", str(e))  # 调试打印
        return {'error': str(e)}, 500

@app.route('/')
def index():
    return 'Hello World'

if __name__ == '__main__':
    app.run(debug=True)